#include <iostream>
#include <vector>
#include <string>
#include <stdlib.h>
#include <math.h>
#include <limits>
#include <fstream>
#include "fitness.hpp"

using namespace std;

vector< vector<bool> > pool;
vector< float > eff; // selection coefficient
vector<string> refType;
vector<string> varType;

// generate poisson random numbers
inline int rpois(float lambda)
{
    float l=exp(-lambda);
    int k=0;
    float p=1.0;
    do
    {
        k++;
        p = p * ((float)rand()) / RAND_MAX;
    }
    while (p > l);
    return k-1;
}

inline float random(float start, float end)
{
    float rndDbl = rand() + 1.0;
    rndDbl = float(start) + ((float)(end-start))*(((float)rndDbl)/((float)RAND_MAX));
    return float(rndDbl);
}

// cross two chromosomes, the actual number of recombinations happened is
// min + poisson(avg-min)
int sporulate(int min, float avg, vector<bool> &id1, vector<bool> &id2);

void randOneCross(vector<bool> &id1, vector<bool> &id2);

// the file named "fn" is no header, tab-delimited file with all the genome-types in it
void loadPool(string fn, int numMarker, int numChr, vector< vector<bool> > &res);

// load the selection coefficient "s".
// positive value means reference state is beneficial.
// negative value means variant state is beneficial.
void loadEffect(string fn, int numMarker, vector<float> &res);

// Get the frequencies of the reference state
void getFreq(vector< vector<bool> > &pop, int numMarker, int numChr, vector<float> &res);

void usage();


int main(int argc, char *argv[])
{
    if(argc != 13) usage();
    string file = argv[1];
    int nMarker = atoi(argv[2]);
    int nHap = atoi(argv[3]);
    int popSize = atoi(argv[4]);
    int minCross = atoi(argv[5]);
    float avgCross = atof(argv[6]);
    string effFile = argv[7];
    int nNormGen = atoi(argv[8]);
    int nSelectGen = atoi(argv[9]);
    int seed1 = atoi(argv[10]);
    int seed2 = atoi(argv[11]);
    string outputFn = argv[12];
//    string outputFn2 = argv[13];

    float sumFitness;
    float thisFitness;
    vector<float> allFitness;
    vector<float> tmpFitness;
    vector< vector<bool> > tmpPool;
    vector<float> selFreq;
    int curInd,tmpInt;

    int curNeed,curGot;
    float x,curFit,rndDbl;

    loadPool(file,nMarker,nHap,pool);
    loadEffect(effFile,nMarker,eff);

    pool.reserve(popSize * 2 + 1);
    allFitness.reserve(popSize);
    tmpPool.reserve(popSize * 2 + 1);
    tmpFitness.reserve(popSize);
    tmpPool = pool;

/*
    // just print where the variant state appears
    for(int i=0;i<nHap;i++) {
        for(int j=0;j<nMarker;j++)
        {
            if(pool[i][j]) cout << i << " x " << j << endl;
        }
    }
*/


    // Generations without selection, with mate and sporulation
    srand(seed1);
    nNormGen ++;
    if(pool.size() == 2) { // pairwise design, get F1 first (F1 is always heterozygous)
        nNormGen --;
        for(int j=0;j<popSize;j++)
        {
            for(int k=0;k<2;k++)
            {
                if((signed)tmpPool.size() > (j * 2 + k) )
                {
                    tmpPool[j * 2 + k] = pool[k];
                }
                else
                {
                    tmpPool.push_back(pool[k]);
                }
            }
            sporulate(minCross,avgCross,tmpPool[j*2],tmpPool[j*2 + 1]);
        }
        pool = tmpPool;
    }

    for(int i=0;i<nNormGen;i++)
    {
        for(int j=0;j<popSize;j++)
        {
            for(int k=0;k<2;k++)
            {
                tmpInt = rand() % pool.size();
                if((signed)tmpPool.size() > (j * 2 + k) )
                {
                    tmpPool[j * 2 + k] = pool[tmpInt];
                }
                else
                {
                    tmpPool.push_back(pool[tmpInt]);
                }
            }
            sporulate(minCross,avgCross,tmpPool[j*2],tmpPool[j*2 + 1]);
        }
        pool = tmpPool;
    }


    // Generations with selection
    srand(seed2);
    for(int i=0;i<nSelectGen;i++) // each generation
    {
        cerr << "Generation "<<i<<endl;
        sumFitness = 0.0;
        for(unsigned int j=0;j<(pool.size()/2);j++) // each individual
        {
            if(allFitness.size() <= (pool.size()/2))
            { // if in the first round, calculate all fitness
                thisFitness = getFitness(pool[j*2],pool[j*2+1],eff);
                allFitness.push_back(thisFitness);
            }
            // if not the first round, no need to recalculate
            sumFitness += allFitness[j];
        }

        //Selection
/*
        // this selection procedure is O(n*log(m) + m)
        // start O(n*log(m) + m)
        double rndDbl;
        for(int j=0;j<popSize;j++)
        {
            rndDbl = random(0,sumFitness);
            curInd = -1;
            while(rndDbl >= 0)
            {
                curInd ++;
                rndDbl = rndDbl - allFitness[curInd];
            }

            if((signed)tmpPool.size() > (j * 2 + 1) )
            {
                tmpPool[j * 2] = pool[curInd * 2];
                tmpPool[j * 2 + 1] = pool[curInd * 2 + 1];
            }
            else
            {
                tmpPool.push_back(pool[curInd * 2]);
                tmpPool.push_back(pool[curInd * 2 + 1]);
            }
            if((signed)tmpFitness.size() > j)
            {
                tmpFitness[j] = allFitness[curInd];
            }
            else
            {
                tmpFitness.push_back( allFitness[curInd] );
            }
        }
        // end O(n*log(m) + m)
*/


        // this selection prcedure is O(m + n)
        // start O(m + n)
        curInd=0;
        curNeed = popSize;
        curGot=0;
        curFit=allFitness[curInd];
        while(curNeed>0) {
            x = sumFitness * ( 1.0 - pow((random(0.0,1.0)),1.0/((float)curNeed) ) );
            sumFitness -= x;
            while(x > curFit) {
                x -= curFit;
                curInd ++;
                if(curInd >= popSize)
                {
                    break;
                }
                curFit = allFitness[curInd];
            }
            if(curInd >= popSize)
            {   // if the fast way did not get enough sample
                // fall back on the slow way
                for(int j=0;j<curNeed;j++)
                {
                    rndDbl = random(0,sumFitness);
                    curInd = -1;
                    while(rndDbl >= 0)
                    {
                        curInd ++;
                        rndDbl = rndDbl - allFitness[curInd];
                    }
                    if((signed)tmpPool.size() > ((curGot +j) * 2 + 1) )
                    {
                        tmpPool[(curGot +j) * 2] = pool[curInd * 2];
                        tmpPool[(curGot +j) * 2 + 1] = pool[curInd * 2 + 1];
                    }
                    else
                    {
                        tmpPool.push_back(pool[curInd * 2]);
                        tmpPool.push_back(pool[curInd * 2 + 1]);
                    }
                    if((signed)tmpFitness.size() > (curGot + j))
                    {
                        tmpFitness[curGot + j] = allFitness[curInd];
                    }
                    else
                    {
                        tmpFitness.push_back( allFitness[curInd] );
                    }
                }
            }
            curFit -= x;

            if((signed)tmpPool.size() > (curGot * 2 + 1) )
            {
                tmpPool[curGot * 2] = pool[curInd * 2];
                tmpPool[curGot * 2 + 1] = pool[curInd * 2 + 1];
            }
            else
            {
                tmpPool.push_back(pool[curInd * 2]);
                tmpPool.push_back(pool[curInd * 2 + 1]);
            }
            if((signed)tmpFitness.size() > curGot)
            {
                tmpFitness[curGot] = allFitness[curInd];
            }
            else
            {
                tmpFitness.push_back( allFitness[curInd] );
            }

            curGot ++;
            curNeed --;
            if(curGot % 100000 == 0) cerr << curGot << endl;
        }
        // end O(m + n)
        pool = tmpPool;
        allFitness = tmpFitness;
    }

    // calculate the frequencies
    getFreq(pool,nMarker ,popSize * 2, selFreq);

    // print results
    ofstream ofs;
    ofs.open(outputFn.c_str());
    for(int i = 0;i<nMarker;i++)
    {
        ofs << selFreq[i] << endl;
    }
    ofs.close();

    /*
    // output detail genome types
    ofs.open(outputFn2.c_str());
    for(int i=0;i<popSize * 2;i++ )
    {
        for(int j=0;j<nMarker;j++)
        {
            if(j == nMarker - 1)
            {
                if(tmpPool[i][j]) {ofs << varType[j] << endl;}
                else{ ofs << refType[j] << endl;}
            }
            else
            {
                if(tmpPool[i][j]) {ofs << varType[j] << '\t';}
                else{ ofs << refType[j] << '\t';}
            }
        }
    }
    ofs.close();
    */


    pool.clear();
    tmpPool.clear();
    allFitness.clear();
    tmpFitness.clear();
    exit(0);
}


void getFreq(vector< vector<bool> > &pop, int numMarker, int numChr, vector<float> &res)
{
    int tmpInt;
    res.clear();
    for(int i=0;i<numMarker;i++)
    {
        tmpInt = 0;
        for(int j=0;j<numChr;j++)
        {
            if(!pop[j][i]) tmpInt ++;
        }
        //cerr << "Marker "<<i<<" : " << tmpInt << endl;
        res.push_back( ((float)tmpInt) / ((float)numChr) );
    }
    return;
}

void loadEffect(string fn, int numMarker, vector<float> &res)
{
    ifstream ifs;
    ifs.open(fn.c_str());
    float tmpDbl;
    res.clear();

    for(int i=0;i<numMarker;i++)
    {
        ifs >> tmpDbl;
        res.push_back(tmpDbl);
    }
    ifs.close();
    return;
}

void loadPool(string fn, int numMarker, int numChr, vector< vector<bool> > &res)
{
    int loadedCnt = 0;
    int markerCnt = 0;
    vector<bool> tmpVecBln;
    tmpVecBln.reserve(numMarker);
    char ch;
    ifstream ifs;
    ifs.open(fn.c_str());
    while(true)
    {
        ch = ifs.peek();
        string strCh(1,ch);
        switch(ch)
        {
        case '\t':
            if(numMarker == markerCnt)
            { // ignore the whole line
                ifs.ignore(numeric_limits<streamsize>::max(),'\n');
                if(loadedCnt % 100000 == 0) {cerr << "Loaded " << loadedCnt << " ; nMarker " << tmpVecBln.size() << endl;}
                res.push_back(tmpVecBln);
                loadedCnt ++;
                tmpVecBln.clear();
                markerCnt = 0;
                if(loadedCnt == numChr)
                {
                    ifs.close();
                    return;
                }
                continue;
            }
            break;
        case -1:
            if(markerCnt == numMarker)
            {
                if(loadedCnt % 100000 == 0) {cerr << "Loaded " << loadedCnt << " ; nMarker " << tmpVecBln.size() << endl;}
                res.push_back(tmpVecBln);
                loadedCnt ++;
                tmpVecBln.clear();
                markerCnt = 0;
                if(loadedCnt == numChr)
                {
                    ifs.close();
                    return;
                }
                else
                {
                    cerr << "Error loading data: the file contains "<< loadedCnt << " useful rows,"<< endl;
                    cerr << "whereas you requested " << numChr << "." << endl;
                    exit(1);
                }
            }
            else
            {
                cerr << "Error loading data: a row in the file contains "<<markerCnt<<" markers,"<<endl;
                cerr << "whereas you requested " << numMarker << "." << endl;
                exit(1);
            }
            break;
        case '\n':
            if(markerCnt == numMarker )
            {
                if(loadedCnt % 100000 == 0) {cerr << "Loaded " << loadedCnt << " ; nMarker " << tmpVecBln.size() << endl;}
                res.push_back(tmpVecBln);
                loadedCnt ++;
                tmpVecBln.clear();
                markerCnt = 0;
                if(loadedCnt == numChr)
                {
                    ifs.close();
                    return;
                }
            }
            else if(markerCnt < numMarker)
            {
                cerr << "Error loading data: a row in the file contains "<<markerCnt<<" markers,"<<endl;
                cerr << "whereas you requested " << numMarker << "." << endl;
                exit(1);
            }
            if(loadedCnt == numChr)
            {
                ifs.close();
                return;
            }
            break;
        default:
            if(loadedCnt == 0)
            {
                refType.push_back(strCh);
                varType.push_back(strCh);
                tmpVecBln.push_back(false);
            }
            else
            {
                if(refType[markerCnt].compare(strCh) == 0 )
                {
                    tmpVecBln.push_back(false);
                }
                else
                {
                    tmpVecBln.push_back(true);
                    varType[markerCnt] = strCh;
                }
            }
            markerCnt ++;
            break;
        }
        ifs.get();
    }
    ifs.close();
    return;
}

int sporulate(int min, float avg, vector<bool> &id1, vector<bool> &id2)
{
    int nCross = min + rpois(avg - min);
    for(int i=0; i<nCross; i++) randOneCross(id1,id2);
    return nCross;
}

void randOneCross(vector<bool> &id1, vector<bool> &id2)
{
    int at = 1 + (rand() % (id1.size() - 1));
    bool tmp;
    for(unsigned int i=at; i<id1.size(); i++)
    {
        tmp = id1[i];
        id1[i] = id2[i];
        id2[i] = tmp;
    }
    return;
}



void usage()
{
    cerr << "Expecting parameters are : (All are needed and in this paricular order)" << endl;
    cerr << "1.markerFile: A tab delimited file listing the markers' states in starting strains" << endl;
    cerr << "    Please use only one character to indicate the state of it." << endl;
    cerr << "    The states of the first row in this file will be regarded as reference state," << endl;
    cerr << "    and all other states are regarded as variants. (So >= 2 states will be reduced to 2)" << endl;
    cerr << "2.nMarker: How many markers to load." << endl;
    cerr << "    Must be smaller than number of columns in markerFile." << endl;
    cerr << "3.nChromosome: How many chromosomes to load." << endl;
    cerr << "    Must be smaller than number of rows in markerFile" << endl;
    cerr << "4.popSize: Population size." << endl;
    cerr << "5.minCross: Minimum number of crossover events per sporulation." << endl;
    cerr << "6.avgCross: Expected number of crossover events per sporulation." << endl;
    cerr << "    The actuall number of crossovers will be (minCross + poisson(avgCross-minCross))" << endl;
    cerr << "7.selCoefFile: A tab delimited file listing the selection coefficient for each marker." << endl;
    cerr << "    A positive value will favored the reference state, and negative value will favored the variant state" << endl;
    cerr << "    E.g. -0.05 means a relative fitness of 1.05 for variant state." << endl;
    cerr << "8.nGenNoSel: Number of generations without selection. During which population size is popSize" << endl;
    cerr << "    Mating and sporulation happend during these generations" << endl;
    cerr << "9.nGenSel: Number of generations with selection. During which population size is popSize." << endl;
    cerr << "    No mating and sporulation during these generations." << endl;
    cerr << "10.randSeed1: Random number seed for generations WITHOUT selection." << endl;
    cerr << "11.randSeed2: Random number seed for generations WITH selection." << endl;
    cerr << "12.freqOutputFile: Output file for the reference state frequencies at the end." << endl;
//    cerr << "13.popOutputFile: Output file for the whole population at the end." << endl;
    exit(0);
}
